# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

import math

import isaaclab.sim as sim_utils
from isaaclab.assets import ArticulationCfg, AssetBaseCfg
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.utils import configclass

from . import mdp

##
# Pre-defined configs
##
# from isaaclab_assets.robots.cartpole import CARTPOLE_CFG  # isort:skip


# -----------------------------------------fanziqi
# import isaaclab.sim as sim_utils
from isaaclab.actuators import DCMotorCfg
# from isaaclab.assets.articulation import ArticulationCfg
YMBOT_E_CFG = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        usd_path="/home/ps/sh_robot/hofa_rl/hofa_rl_isaaclab/source/hofa_rl_isaaclab/hofa_rl_isaaclab/assets/ymbot_e/ymbot_e.usd",
        activate_contact_sensors=True,
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=False,
            retain_accelerations=False, 
            linear_damping=0.0,
            angular_damping=0.0,
            max_linear_velocity=1000.0,
            max_angular_velocity=1000.0,
            max_depenetration_velocity=1.0, 
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False, 
            solver_position_iteration_count=4, 
            solver_velocity_iteration_count=0,
            fix_root_link=True,
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        pos=(0.0, 0.0, 1.8),
        joint_pos={
            # ".*L_hip_joint": 0.0,
            # ".*R_hip_joint": -0.0,
            # "F.*_thigh_joint": 0.8,
            # "R.*_thigh_joint": 0.8,
            # ".*_calf_joint": -1.5,
            ".*": 0.6 # !!!
        },
        joint_vel={".*": 0.0},
    ),
    soft_joint_pos_limit_factor=0.9,
    actuators={
        "legs": DCMotorCfg(
            # joint_names_expr=[".*_Joint\\d*"],
            joint_names_expr=["Neck_Joint2", "Torso_Joint2"],
            effort_limit=200.0, #33.5,
            saturation_effort=200.0, #33.5,
            velocity_limit=21.0,
            stiffness=20.0,
            damping=0.5,
            friction=0.0,
        ),
    },
)
"""Configuration of YMBOT_E using DC motor."""


##
# Scene definition
##


@configclass
class YmbotEPidSceneCfg(InteractiveSceneCfg):
    """Configuration for a cart-pole scene."""

    # ground plane
    ground = AssetBaseCfg(
        prim_path="/World/ground",
        spawn=sim_utils.GroundPlaneCfg(size=(100.0, 100.0)),
    )

    # robot
    robot: ArticulationCfg = YMBOT_E_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
    # robot: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")

    # lights
    dome_light = AssetBaseCfg(
        prim_path="/World/DomeLight",
        spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0),
    )


##
# MDP settings
##

# @configclass
# class ActionsCfg:
#     """Action specifications for the MDP."""

#     joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=["slider_to_cart"], scale=100.0)
#----------------------------------------------------------------------------------------------
# from .mdp.pd_joint_actions import PDJointActionCfg
# @configclass
# class ActionsCfg:
#     """MDP 的动作规范。"""

#     joint_position_to_effort_new = PDJointActionCfg(
#         asset_name="robot",
#         joint_names=["slider_to_cart"],
#         scale=1.0,
#         kp=10.0,  # 比例增益
#         kd=1.0    # 微分增益
#     )
#----------------------------------------------------------------------------------------------
# from .mdp.adaptive_pd_joint_actions import AdaptivePDJointActionCfg
# @configclass
# class ActionsCfg:
#     """MDP 的动作规范。"""

#     adaptive_pd_joint_actions = AdaptivePDJointActionCfg(
#         asset_name="robot",
#         joint_names=["slider_to_cart"],
#         initial_kp=10.0,
#         initial_kd=1.0,
#         target_pos=[0.0]  # 使用列表表示目标位置
#     )
#----------------------------------------------------------------------------------------------
from .mdp.adaptive_pd_joint_actions import AdaptivePDJointActionCfg
@configclass
class ActionsCfg:
    """MDP 的动作规范。"""

    adaptive_pd_joint1 = AdaptivePDJointActionCfg(
        asset_name="robot",
        joint_names=["Neck_Joint2"],
        initial_kp=10.0,
        initial_kd=1.0,
        target_pos=[0.0]  # 使用列表表示目标位置
    )
    adaptive_pd_joint2 = AdaptivePDJointActionCfg(
        asset_name="robot",
        joint_names=["Torso_Joint2"],
        initial_kp=10.0,
        initial_kd=1.0,
        target_pos=[0.0]  # 使用列表表示目标位置
    )
#----------------------------------------------------------------------------------------------


@configclass
class ObservationsCfg:
    """Observation specifications for the MDP."""

    @configclass
    class PolicyCfg(ObsGroup):
        """Observations for policy group."""

        # observation terms (order preserved)
        joint_pos_rel = ObsTerm(func=mdp.joint_pos_rel)
        joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel)

        def __post_init__(self) -> None:
            self.enable_corruption = False
            self.concatenate_terms = True

    # observation groups
    policy: PolicyCfg = PolicyCfg()


@configclass
class EventCfg:
    """Configuration for events."""

    # # reset
    # reset_cart_position = EventTerm(
    #     func=mdp.reset_joints_by_offset,
    #     mode="reset",
    #     params={
    #         "asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]),
    #         "position_range": (-1.0, 1.0),
    #         "velocity_range": (-0.5, 0.5),
    #     },
    # )

    # reset_pole_position = EventTerm(
    #     func=mdp.reset_joints_by_offset,
    #     mode="reset",
    #     params={
    #         "asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]),
    #         "position_range": (-0.25 * math.pi, 0.25 * math.pi),
    #         "velocity_range": (-0.25 * math.pi, 0.25 * math.pi),
    #     },
    # )
    # reset
    reset_joint1_position = EventTerm(
        func=mdp.reset_joints_by_offset,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=["Neck_Joint2"]),
            "position_range": (-0.25 * math.pi, 0.25 * math.pi),
            "velocity_range": (-0.25 * math.pi, 0.25 * math.pi),
        },
    )

    reset_joint2_position = EventTerm(
        func=mdp.reset_joints_by_offset,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=["Torso_Joint2"]),
            "position_range": (-0.25 * math.pi, 0.25 * math.pi),
            "velocity_range": (-0.25 * math.pi, 0.25 * math.pi),
        },
    )


@configclass
class RewardsCfg:
    """Reward terms for the MDP."""

    # # (1) Constant running reward
    # alive = RewTerm(func=mdp.is_alive, weight=1.0)
    # # (2) Failure penalty
    # terminating = RewTerm(func=mdp.is_terminated, weight=-2.0)
    # # (3) Primary task: keep pole upright
    # pole_pos = RewTerm(
    #     func=mdp.joint_pos_target_l2,
    #     weight=-1.0,
    #     params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), "target": 0.0},
    # )
    # # (4) Shaping tasks: lower cart velocity
    # cart_vel = RewTerm(
    #     func=mdp.joint_vel_l1,
    #     weight=-0.01,
    #     params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"])},
    # )
    # # (5) Shaping tasks: lower pole angular velocity
    # pole_vel = RewTerm(
    #     func=mdp.joint_vel_l1,
    #     weight=-0.005,
    #     params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"])},
    # )

    # (1) Constant running reward
    alive = RewTerm(func=mdp.is_alive, weight=1.0)
    # (2) Failure penalty
    terminating = RewTerm(func=mdp.is_terminated, weight=-2.0)
    # (3) Primary task: keep pole upright
    joint1 = RewTerm(
        func=mdp.joint_pos_target_l2,
        weight=-10.0,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Neck_Joint2"]), "target": 0.0},
    )
    joint2 = RewTerm(
        func=mdp.joint_pos_target_l2,
        weight=-10.0,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Torso_Joint2"]), "target": 0.0},
    )
    # (4) Shaping tasks:  lower joint angular velocity
    joint1_vel = RewTerm(
        func=mdp.joint_vel_l1,
        weight=-0.01,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Neck_Joint2"])},
    )
    joint2_vel = RewTerm(
        func=mdp.joint_vel_l1,
        weight=-0.01,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Torso_Joint2"])},
    )


@configclass
class TerminationsCfg:
    """Termination terms for the MDP."""

    # # (1) Time out
    # time_out = DoneTerm(func=mdp.time_out, time_out=True)
    # # (2) Cart out of bounds
    # cart_out_of_bounds = DoneTerm(
    #     func=mdp.joint_pos_out_of_manual_limit,
    #     params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), "bounds": (-3.0, 3.0)},
    # )


    # (1) Time out
    time_out = DoneTerm(func=mdp.time_out, time_out=True)
    # (2) joint out of bounds
    joint1_out_of_bounds = DoneTerm(
        func=mdp.joint_pos_out_of_manual_limit,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Neck_Joint2"]), "bounds": (-math.pi/3 , math.pi/3)},
    )
    joint2_out_of_bounds = DoneTerm(
        func=mdp.joint_pos_out_of_manual_limit,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=["Torso_Joint2"]), "bounds": (-math.pi/3 , math.pi/3)},
    )



##
# Environment configuration
##


@configclass
class YmbotEPidEnvCfg(ManagerBasedRLEnvCfg):
    # Scene settings
    scene: YmbotEPidSceneCfg = YmbotEPidSceneCfg(num_envs=4096, env_spacing=4.0)
    # Basic settings
    observations: ObservationsCfg = ObservationsCfg()
    actions: ActionsCfg = ActionsCfg()
    events: EventCfg = EventCfg()
    # MDP settings
    rewards: RewardsCfg = RewardsCfg()
    terminations: TerminationsCfg = TerminationsCfg()

    # Post initialization
    def __post_init__(self) -> None:
        """Post initialization."""
        # general settings
        self.decimation = 2
        self.episode_length_s = 5
        # viewer settings
        self.viewer.eye = (8.0, 0.0, 5.0)
        # simulation settings
        self.sim.dt = 1 / 120
        self.sim.render_interval = self.decimation


    # def __post_init__(self) -> None:
    #     """Post initialization."""
    #     # general settings
    #     self.decimation = 10
    #     self.episode_length_s = 5
    #     # viewer settings
    #     self.viewer.eye = (8.0, 0.0, 5.0)
    #     # simulation settings
    #     self.sim.dt = 1 / 1000
    #     self.sim.render_interval = self.decimation